from model import CUnet, GaussianDiffusion
from trainer import Trainer
from omegaconf import OmegaConf
import argparse
import os
import torch.multiprocessing as mp
import torch

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    
    parser.add_argument("--config", type=str, default='config/cfg.yaml')
    parser.add_argument("--gpu", type=int, nargs='+', default=0)
    parser.add_argument("--debug", action="store_true")
    args = parser.parse_args()
    
    config = OmegaConf.load(args.config)
    
    gpu = args.gpu
    
    # mp.set_start_method('spawn', force=True)
    
    # os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
    
    torch.cuda.set_device(1)
    
    unet_kwargs = config['Unet']
    diffusion_kwargs = config['GaussianDiffusion']
    
    model = CUnet(**unet_kwargs)
    diffusion = GaussianDiffusion(model, **diffusion_kwargs)
    
    trainer_kwargs = config['Trainer']
    trainer = Trainer(diffusion, **trainer_kwargs, debug=args.debug)
    
    if config.get('resume', None):
        print(f'Resume training from {config["resume"]} milestone')
        trainer.load(config['resume'])
    
    trainer.train()